import SPTAG
import numpy as np

# 假设 SIFT1M fbin 文件路径
fbin_path = '/home/gary/Code/DiskANN/build/data/bigann/bigann_base.bbin'

# 读取fbin文件
def read_bbin(file_path):
    with open(file_path, 'rb') as f:
        num_vectors = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取向量数量
        dim = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取维度
        print("num_vectors:",num_vectors,"dim:",dim)
        print("Reading...")
        data = np.fromfile(f, dtype=np.uint8).reshape(num_vectors, dim)  # 读取所有向量并重塑为矩阵
        print("Read done")
    return data

# 加载SIFT1M数据
sift1b_data = read_bbin(fbin_path)

# 创建SPTAG索引实例
index_builder = SPTAG.AnnIndex('KDT', 'UInt8', sift1b_data.shape[1])
index_builder.SetBuildParam('NumberOfThreads', '46', "Index")  # 设置构建时使用的线程数
index_builder.SetBuildParam('Samples', '100000', "Index")       # 设置采样数量
index_builder.SetBuildParam('HashTableExponent', '4', "Index")

# 插入数据
print("Training...")
index_builder.Build(sift1b_data, sift1b_data.shape[0], False)
print("Training Done...")

# 保存索引
index_builder.Save('sift1b_index_py_full_240814')

print("索引构建完成并已保存。")